import cv2
import numpy as np
import json
import os
from typing import List, Dict, Tuple

class RegionSelector:
    """交互式区域选择器，支持UVC/RTSP配置化拉流"""
    
    def __init__(self, config_file="config.json"):
        self.config_file = config_file
        self.regions = []
        self.current_region = None
        self.dragging = False
        self.start_point = None
        self.current_name = ""
        self.window_name = "[Area Selector] Save(S) Clear(C) Exit(Q)"
        
        # 拉流配置
        self.source_type = "uvc"  # 默认拉流方式
        self.camera_id = 0        # UVC默认设备ID
        self.rtsp_url = ""        # RTSP流地址
        self.camera_width = 1280
        self.camera_height = 720
        self._load_stream_config()

    def _load_stream_config(self) -> None:
        """加载视频流配置"""
        try:
            if os.path.exists(self.config_file):
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    if "camera" in config:
                        camera_config = config["camera"]
                        self.source_type = camera_config.get("source_type", "uvc").lower()
                        
                        # 加载UVC配置
                        uvc_config = camera_config.get("uvc", {})
                        self.camera_id = uvc_config.get("id", 0)
                        
                        # 加载RTSP配置
                        rtsp_config = camera_config.get("rtsp", {})
                        self.rtsp_url = rtsp_config.get("url", "")
                        
                        # 加载分辨率配置（优先使用当前拉流方式的配置）
                        if self.source_type == "rtsp":
                            self.camera_width = rtsp_config.get("width", 1280)
                            self.camera_height = rtsp_config.get("height", 720)
                        else:
                            self.camera_width = uvc_config.get("width", 1280)
                            self.camera_height = uvc_config.get("height", 720)
                        
                        print(f"已加载 {self.source_type.upper()} 配置 | 分辨率: {self.camera_width}x{self.camera_height}")
        except Exception as e:
            print(f"加载视频流配置失败: {e}")
    
    def load_config(self) -> None:
        """从配置文件加载已有区域"""
        if os.path.exists(self.config_file):
            try:
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    if "detection" in config and "regions" in config["detection"]:
                        self.regions = []
                        for region in config["detection"]["regions"]:
                            self.regions.append({
                                "x": region["x"],
                                "y": region["y"],
                                "width": region["width"],
                                "height": region["height"],
                                "name": region["name"]
                            })
                print(f"已从 {self.config_file} 加载 {len(self.regions)} 个区域")
            except Exception as e:
                print(f"加载配置失败: {e}")
    
    def save_config(self) -> bool:
        """保存区域到配置文件"""
        try:
            # 先加载现有配置
            config = {}
            if os.path.exists(self.config_file):
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    config = json.load(f)
            
            # 更新区域配置
            if "detection" not in config:
                config["detection"] = {}
            config["detection"]["regions"] = self.regions
            
            # 保存配置
            with open(self.config_file, 'w', encoding='utf-8') as f:
                json.dump(config, f, ensure_ascii=False, indent=4)
            print(f"已保存 {len(self.regions)} 个区域到 {self.config_file}")
            return True
        except Exception as e:
            print(f"保存配置失败: {e}")
            return False
    
    def mouse_callback(self, event, x, y, flags, param) -> None:
        """鼠标事件回调函数"""
        if event == cv2.EVENT_LBUTTONDOWN:
            # 开始拖拽
            self.dragging = True
            self.start_point = (x, y)
        elif event == cv2.EVENT_MOUSEMOVE and self.dragging:
            # 拖拽中
            self.current_region = {
                "x": min(self.start_point[0], x),
                "y": min(self.start_point[1], y),
                "width": abs(x - self.start_point[0]),
                "height": abs(y - self.start_point[1]),
                "name": f"area{len(self.regions) + 1}" if not self.current_name else self.current_name
            }
        elif event == cv2.EVENT_LBUTTONUP:
            # 结束拖拽
            self.dragging = False
            if self.current_region and self.current_region["width"] > 10 and self.current_region["height"] > 10:
                # 使用默认名称，避免输入阻塞
                self.current_region["name"] = f"area{len(self.regions) + 1}"
                self.regions.append(self.current_region)
                print(f"已添加区域: {self.current_region['name']} "
                      f"({self.current_region['x']}, {self.current_region['y']}, "
                      f"{self.current_region['width']}, {self.current_region['height']})")
                self.current_region = None
    
    def draw_regions(self, frame) -> np.ndarray:
        """在图像上绘制所有区域"""
        # 绘制已保存的区域
        for i, region in enumerate(self.regions):
            x, y = region["x"], region["y"]
            w, h = region["width"], region["height"]
            name = region["name"]
            cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
            cv2.putText(frame, f"{i+1}: {name}", (x, y-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        # 绘制当前正在选择的区域
        if self.current_region:
            x, y = self.current_region["x"], self.current_region["y"]
            w, h = self.current_region["width"], self.current_region["height"]
            cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 0, 255), 2)
        
        return frame
    
    def run(self) -> None:
        """运行区域选择器"""
        # 加载已有配置
        self.load_config()
        
        # 创建视频捕获对象
        cap = None
        if self.source_type == "rtsp":
            if not self.rtsp_url:
                print("错误: RTSP URL未配置")
                return
            try:
                # 使用FFMPEG后端提高兼容性
                cap = cv2.VideoCapture(self.rtsp_url, cv2.CAP_FFMPEG)
                print(f"正在尝试连接 RTSP 流: {self.rtsp_url}")
            except Exception as e:
                print(f"RTSP连接异常: {e}")
                return
        else:
            try:
                # 使用DShow后端提高UVC兼容性
                cap = cv2.VideoCapture(self.camera_id, cv2.CAP_DSHOW)
                print(f"正在尝试打开 UVC 摄像头 ID: {self.camera_id}")
            except Exception as e:
                print(f"UVC摄像头打开异常: {e}")
                return

        # 检查是否成功打开
        if not cap or not cap.isOpened():
            print(f"无法打开 {'RTSP' if self.source_type == 'rtsp' else 'UVC'} 视频源")
            return

        # 设置分辨率（主要对UVC有效）
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.camera_width)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.camera_height)
        
        # 创建窗口并设置鼠标回调
        cv2.namedWindow(self.window_name)
        
        while True:
            # 捕获一帧图像
            ret, frame = cap.read()
            if not ret:
                print("无法捕获图像")
                break
            
            # 绘制区域
            display_frame = frame.copy()
            display_frame = self.draw_regions(display_frame)
            
            # 显示操作说明
            cv2.putText(display_frame, "Drag to select an area: Press 'S' to save, 'C' to clear, 'Q' to exit", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
            
            # 显示图像
            cv2.imshow(self.window_name, display_frame)
            cv2.setMouseCallback(self.window_name, self.mouse_callback, display_frame)
            
            # 处理键盘事件
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                # 退出
                break
            elif key == ord('s'):
                # 保存配置
                self.save_config()
            elif key == ord('c'):
                # 清除所有区域
                self.regions = []
                self.current_region = None
                print("已清除所有区域")
                # 清除后立即保存配置，确保下次启动时不会加载旧区域
                self.save_config()
        
        # 释放资源
        cap.release()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    # 从配置文件自动获取配置，无需手动指定摄像头ID
    selector = RegionSelector(config_file="config.json")
    selector.run()